import numpy as np


def expand_list(input, batchsize):
    coors = list()
    coors.append(input)
    for i in range(batchsize - 1):
        coors.append(np.zeros_like(input))
    return coors

def expand_input(input):
    coors = []
    for i, coor in enumerate(input):
        coor_pad = np.pad(
            coor, ((0, 0), (1, 0)), mode="constant", constant_values=i
        )
        coors.append(coor_pad)
    coors = np.concatenate(coors, axis=0)
    return coors


def train_collate(points, voxels, shape, num_points, num_voxels, coordinates,ref_from_car, car_from_global, 
            hm0, hm1, hm2, hm3, hm4, hm5, anno_box, ind, mask, cat, batch_info):
    batchsize = len(points)
    voxels = expand_list(np.concatenate(voxels, axis=0), batchsize)
    num_points = expand_list(np.concatenate(num_points, axis=0), batchsize)
    points = expand_list(expand_input(points), batchsize)
    coordinates = expand_list(expand_input(coordinates), batchsize)
    # print(len(points), points[0].shape)
    # voxels = np.concatenate(voxels, axis=0)
    # num_points =np.concatenate(num_points, axis=0)
    # points = expand_input(points)
    # coordinates = expand_input(coordinates)
    return points, voxels, shape, num_points, num_voxels, coordinates,ref_from_car, car_from_global,\
    hm0, hm1, hm2, hm3, hm4, hm5, anno_box, ind, mask, cat #token, 


def eval_collate(token, points, voxels, shape, num_points, num_voxels, coordinates,ref_from_car, car_from_global, batch_info):
    batchsize = len(token)
    # print(batchsize)
    voxels = expand_list(np.concatenate(voxels, axis=0), batchsize)
    num_points = expand_list(np.concatenate(num_points, axis=0), batchsize)
    points = expand_list(expand_input(points), batchsize)
    coordinates = expand_list(expand_input(coordinates), batchsize)

    return token, points, voxels, shape, num_points, num_voxels, coordinates,ref_from_car, car_from_global